import csv
import argparse

# 设置命令行参数解析
parser = argparse.ArgumentParser(description='生成机场航线测试数据')
parser.add_argument('--max_edges', type=int, default=200, 
                    help='最大边数量限制（默认：200）')
args = parser.parse_args()

# 读取机场信息
airports = {}
with open('airports.dat', 'r', encoding='utf-8') as f:
    reader = csv.reader(f)
    for row in reader:
        try:
            airport_id = row[0].strip()
            lon = float(row[7].strip())
            lat = float(row[6].strip())
            airports[airport_id] = (lon, lat)
        except (IndexError, ValueError):
            continue

# 读取航线信息并筛选有效边
edges = []
used_node_ids = set()  # 🔧 新增：用于记录实际使用的 node id

with open('routes.dat', 'r', encoding='utf-8') as f:
    reader = csv.reader(f)
    for row in reader:
        try:
            source_id = row[3].strip()
            dest_id = row[5].strip()
            if source_id in airports and dest_id in airports:
                edges.append((source_id, dest_id))
                used_node_ids.update([source_id, dest_id])  # 🔧 记录参与的节点
                if len(edges) >= args.max_edges:  # 🔧 使用命令行参数控制边数量
                    break
        except IndexError:
            continue

# 写入 edges.csv（不变）
with open('edges.csv', 'w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    writer.writerow(['from', 'to'])
    writer.writerows(edges)

# 写入 nodes.csv（只保留参与的 node）🔧
with open('nodes.csv', 'w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    writer.writerow(['id', 'lon', 'lat'])
    for aid in used_node_ids:
        lon, lat = airports[aid]
        writer.writerow([aid, lon, lat])

print(f"✅ 处理完成：{len(used_node_ids)} 个节点，{len(edges)} 条边")